import cv2
import matplotlib.image as mplimg
import matplotlib.pyplot as plt
import numpy as np

blur_kernel_size = 5  # Gaussian blur kernel size
canny_low_threshold = 50  # Canny edge detection low threshold
canny_high_threshold = 150  # Canny edge detection high threshold

# Hough transform parameters
rho = 1
theta = np.pi / 180
threshold = 15
min_line_length = 40
max_line_gap = 20


def roi_mask(img, vertices):
    """
    划分出区域
    :param img:
    :param vertices:
    :return:
    """
    mask = np.zeros_like(img)
    mask_color = 255
    cv2.fillPoly(mask, vertices, mask_color)
    # 绘制多边形 在mask上绘制  vertices顶点，指的是车道线多边形区域的顶点
    masked_img = cv2.bitwise_and(img, mask)
    # 按位与   返回的是八位二进制，作为霍夫变换的输入
    return masked_img


def hough_lines(img, rho, theta, threshold,
                min_line_len, max_line_gap):
    """
    霍夫变换 找出直线
    :param img: 需要处理的图像
    :param rho: 直线的精度
    :param theta:
    :param threshold: 域值，超过阈值的线才能被返回
    :param min_line_len: 线的最小距离
    :param max_line_gap: 线的最大距离
    :return: 霍夫变换后的直线
    """
    lines = cv2.HoughLinesP(img, rho, theta, threshold, np.array([]),
                            minLineLength=min_line_len,
                            maxLineGap=max_line_gap)
    line_img = np.zeros((img.shape[0], img.shape[1], 3), dtype=np.uint8)  # 全零的数组
    draw_lanes(line_img, lines)
    return line_img


def clean_lines(lines, threshold):
    slope = [(y2 - y1) / (x2 - x1) for line in lines for x1, y1, x2, y2 in line]  # 斜率
    while len(lines) > 0:
        mean = np.mean(slope)  # 求均值
        diff = [abs(s - mean) for s in slope]  # 斜率和均值的差的绝对值
        idx = np.argmax(diff)  # 求某一维度中数值最大的索引值
        if diff[idx] > threshold:
            slope.pop(idx)  # 删除最后一个元素
            lines.pop(idx)
        else:
            break


def draw_lanes(img, lines, color=[255, 0, 0], thickness=8):
    """
    划线，在ROI区域  lines为上面红色线
    :param img:
    :param lines:
    :param color:
    :param thickness:
    :return:
    """
    left_lines, right_lines = [], []
    for line in lines:
        for x1, y1, x2, y2 in line:
            k = (y2 - y1) / (x2 - x1)
            if k < 0:
                left_lines.append(line)  # 增加元素
            else:
                right_lines.append(line)

    if (len(left_lines) <= 0 or len(right_lines) <= 0):
        return img  # 判断斜率并返回
    clean_lines(left_lines, 0.1)
    clean_lines(right_lines, 0.1)
    left_points = [(x1, y1) for line in left_lines for x1, y1, x2, y2 in line]
    left_points = left_points + [(x2, y2) for line in left_lines for x1, y1, x2, y2 in line]
    right_points = [(x1, y1) for line in right_lines for x1, y1, x2, y2 in line]
    right_points = right_points + [(x2, y2) for line in right_lines for x1, y1, x2, y2 in line]
    left_vtx = calc_lane_vertices(left_points, 325, img.shape[0])
    right_vtx = calc_lane_vertices(right_points, 325, img.shape[0])
    cv2.line(img, left_vtx[0], left_vtx[1], color, thickness)
    # 参数含义：划线函数， ROI区域，直线起点，终点，颜色，厚度
    cv2.line(img, right_vtx[0], right_vtx[1], color, thickness)


def calc_lane_vertices(point_list, ymin, ymax):
    x = [p[0] for p in point_list]
    y = [p[1] for p in point_list]
    fit = np.polyfit(y, x, 1)  # 最小二乘法拟合
    fit_fn = np.poly1d(fit)  # 生成拟合的式子
    xmin = int(fit_fn(ymin))
    xmax = int(fit_fn(ymax))
    return [(xmin, ymin), (xmax, ymax)]


def main():
    # 读取图片
    image = cv2.imread('images/road.jpg')
    # 转化为灰度图
    gray = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY)
    # 高斯滤波
    blur_gray = cv2.GaussianBlur(gray, (blur_kernel_size, blur_kernel_size), 0, 0)
    # 霍夫变换
    edges = cv2.Canny(blur_gray, canny_low_threshold, canny_high_threshold)

    # 划分roi区
    roi_vtx = np.array([[(0, edges.shape[0]), (460, 325),
                         (520, 325), (edges.shape[1], edges.shape[0])]])  # shape[0]=540   [1]=940

    # 车道线的位置，这个应该是根据车道线的在图像中的具体位置划分的
    roi_edges = roi_mask(edges, roi_vtx)  # 结果返回为二进制

    line_img = hough_lines(roi_edges, rho, theta, threshold,
                           min_line_length, max_line_gap)

    # 叠加图像
    res_img = cv2.addWeighted(image, 0.8, line_img, 1, 0)

    res_image = res_img[..., ::-1]
    plt.imshow(res_image)
    plt.savefig('res_image.jpg')
    plt.show()


if __name__ == '__main__':
    main()
